#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
用于读取多个traceSeq生成的h5文件并绘制相关图表的脚本
支持读取多个文件，按距离排列traces，进行统一的FV变换
"""

import argparse
import os
import sys
import numpy as np
from pylab import plt
import pandas as pd
from obspy.core.utcdatetime import UTCDateTime
import glob

# 导入自定义工具
from utils.h5data import read_h5_data
from utils.trace import filtfilt, stack_traces_in_bins
from utils.fv import do_fv
from utils.plot import plot_traces, plot_raw_data
from utils.loc import load_loc, get_distance,sort_data_by_distance
from utils.math import norm

# 命令行参数解析
parser = argparse.ArgumentParser(description='Plot traceSeq h5 data from multiple files')
parser.add_argument('-input', default='', help='input h5 file pattern')
parser.add_argument('-figroot', default='figures/6.traceSeq.figures', help='root to save figs')
parser.add_argument('-emarker', default='', help='marker name')
parser.add_argument('-noD1', action='store_true', help='method: discard S2N events')
parser.add_argument('-noD2', action='store_true', help='method: discard N2S events')
parser.add_argument('-noD3', action='store_true', help='method: discard mixed events')
args = parser.parse_args()


INPUT_PATTERN = args.input
EMARKER = args.emarker

FIG_ROOT = args.figroot
USE_D1_EVENTS = not args.noD1
USE_D2_EVENTS = not args.noD2
USE_D3_EVENTS = not args.noD3

fs, fe = 5, 40
# VLIM = [100,1400]
VLIM = [100,4000]
TLIM = [-2, 2]
NORM = True
PLOT_WIGGLE = True 
STACK = True
FOLDER = False

MARKER = f'F{fs:.1f}.{fe:.1f}.V{VLIM[0]:.0f}.{VLIM[1]:.0f}.N.{int(NORM)}.S.{int(STACK)}.F.{int(FOLDER)}'
figname = f'{FIG_ROOT}/{EMARKER}.{MARKER}.png'

# 获取所有匹配的文件
if not INPUT_PATTERN:
    print("请提供文件匹配模式，例如: data/2406.HSR/events/P192.2406.Z.traces.P*.h5")
    sys.exit(1)

input_files = glob.glob(INPUT_PATTERN)
input_files.sort()
if not input_files:
    print(f"未找到匹配的文件: {INPUT_PATTERN}")
    sys.exit(1)

print(f"找到 {len(input_files)} 个文件")

# 读取第一个文件获取基本参数
datasets, args_infile = read_h5_data(input_files[0], keys=['MARKER','EMARKER'], group_name='metadata', read_attrs=True)
DATE = args_infile['date']
EMARKER = f'{datasets[1].decode()}.{datasets[0].decode()}' if not EMARKER else EMARKER

# 存储所有数据
all_ref_traces = []
all_x = []
all_corr_stats = []
train_pos = []
# 读取所有文件的数据
print('Reading data from all files...')
for i, input_file in enumerate(input_files):
    print(f'  Reading {i+1}/{len(input_files)}: {os.path.basename(input_file)}')
    
    # 读取h5文件中的数据
    try:
        data_dict = read_h5_data(input_file, keys=['t','ref', 'S', 'R', 'x'])
        _, args_infile1 = read_h5_data(input_file, keys=['MARKER'], group_name='metadata', read_attrs=True)
        tn, ref_traces, S_data, R, x = data_dict
        vinfo = read_h5_data(input_file, keys=['info'])[0]
        # 解码字符串数组
        S = S_data.decode()
        ns,nt = ref_traces.shape
        # idx_events = np.where(vinfo==2)[0]
        print()
                              
        for j in range(ns):
            name_j = R[j].decode()
            if 'E' not in name_j and abs(x[j])<3000:
                all_corr_stats.append(S+'/'+name_j)
                all_ref_traces.append(ref_traces[j,:])
                # all_ref_traces.append(traces[idx_events,j,:].sum(axis=0))
                all_x.append(x[j])
        train_pos.append(args_infile1['Tpos'])
            
    except Exception as e:
        print(f"  跳过文件 {input_file}，错误: {e}")

# 合并所有数据
if not all_ref_traces:
    print("没有成功读取任何数据")
    sys.exit(1)

ref_traces_combined = np.array(all_ref_traces)
x_combined = np.array(all_x)
# 按距离排序
sort_idx = np.argsort(x_combined)
x_combined = x_combined[sort_idx]
ref_traces_combined = ref_traces_combined[sort_idx,:]
names = [all_corr_stats[i] for i in sort_idx]

ns = len(all_corr_stats)

print(f'总共有 {ns} 个 traces')

if STACK:
    # 每32米范围内记录叠加生成新记录
    ref_traces_combined, x_combined, names = stack_traces_in_bins(ref_traces_combined, x_combined)
if FOLDER:
    ref_traces_combined = ref_traces_combined+ref_traces_combined[:,::-1]

# 数据处理
ref_traces = filtfilt(ref_traces_combined, tn[1]-tn[0], [fs, fe], order=4)
ref_traces = norm(ref_traces, ONE_AXIS=True)
x = np.array(x_combined)
ns = len(x)
print(ref_traces.shape)

# 绘制参考道数据
print('Plotting reference traces...')



vs = np.linspace(VLIM[0],VLIM[1],800)
FV, f,_ = do_fv(ref_traces,vs,x=np.array(x), sx=0,dt=tn[1]-tn[0], NORM=True, PRINT=False, fs=fs, fe=fe)
FV = np.abs(FV)
FV = FV/FV.max()

# 绘制图形
fig, ax1, ax2 = plot_raw_data(ref_traces, x, tn, fs=fs, fe=fe, 
                              VLIM=VLIM, PLOT_WIGGLE=PLOT_WIGGLE, SCALE=40, FV_NORM=NORM)

if PLOT_WIGGLE:
    # for j in range(30):  # 限制标签数量避免过于拥挤
    #     ax1.text(-1, x[j], all_corr_stats[j], fontsize=6)
    ax1.set_ylim([x.min()-100, x.max()+100])
else:
    ax1.set_ylim([x.min(), x.max()])

ax1.set_title(f'{EMARKER}')
ax1.set_xlim(TLIM)

# 频散曲线
try:
    disp = np.loadtxt('figures/6.traceSeq.figures/disp.txt', delimiter=',')
    # ax2.plot(disp[:,0], disp[:,1], 'yellow', lw=1)
    # ax2.plot(disp[:,0], -disp[:,1], 'yellow', lw=1)
except:
    print("未找到 disp.txt 文件，跳过频散曲线绘制")

D = 10 # unit: m
FMAX = fe # unit: Hz
FMIN = fs # unit: Hz
CS = 1400 # unit: m/s
CP = 3500 # unit: m/s

CMAX = 10000
from lambwaves import Lamb
alum = Lamb(thickness=D*1000,  #unit: mm 
            nmodes_sym=2,     # mode
            nmodes_antisym=2, # mode
            fd_max=FMAX*D,     # fd/thickness 是 kHz*mm
            vp_max=CMAX,     # unit: m/s
            c_L=CP,         # unit: m/s
            c_S=CS)         # unit: m/s
freqs = np.linspace(FMIN, FMAX, 300)
c_a0 = alum.vp_antisym['A0'](freqs*D)
# ax2.plot(freqs, c_a0, 'k-', lw=1)
# ax2.text(freqs[0], c_a0[0],'A0')


D = 1 # unit: m
FMAX = fe # unit: Hz
FMIN = fs # unit: Hz
CS = 450 # unit: m/s
CP = 1000 # unit: m/s

CMAX = 10000
from lambwaves import Lamb
alum = Lamb(thickness=D*1000,  #unit: mm 
            nmodes_sym=2,     # mode
            nmodes_antisym=2, # mode
            fd_max=FMAX*D+2,     # fd/thickness 是 kHz*mm
            vp_max=CMAX,     # unit: m/s
            c_L=CP,         # unit: m/s
            c_S=CS)         # unit: m/s
freqs = np.linspace(FMIN, FMAX, 300)
c_a0 = alum.vp_sym['S0'](freqs*D)
# ax2.plot(freqs, c_a0, 'k--', lw=1)
# ax2.text(freqs[0], c_a0[0],'S0')

# # 参考速度
# ax1.plot([-5, 5], [-5*3000, 5*3000], 'k-', lw=1)
# ax1.plot([-5, 5], [-5*340, 5*340], 'k--', lw=1)

# # 添加特定频率线
# f0 = 4.4
# for j in [1, 4, 9]:
#     ax2.plot([j*f0, j*f0], VLIM, color='yellow', lw=1)

ax2.set_ylim(VLIM)
fig.tight_layout()

# 保存图像
fig.savefig(figname, dpi=300)
print(f'Saved {figname}')

print('Plotting completed.')